# data_loader.py
import numpy as np
import os

from scipy.interpolate import interp1d
from torch import nn

from auth_config import LOS_FP_PATH, NLOS_FP_PATH, DEVICE_LIST, DL_FP_PATH, LOS2_FP_PATH, NLOS2_FP_PATH,  \
    RAW_CSI_FEAT_DIM, DL2_FP_PATH, MOV_FP_PATH, MOV2_FP_PATH



# 在tools.py中添加
import scipy.io as sio

def load_raw_csi(device_name, scenario="LOS"):
    """
    加载原始CSI复数数据（.mat格式）
    :param device_name: 设备名（如"DUT01"）
    :param scenario: 场景（"LOS"或"NLOS"）
    :return: 处理后的特征数组 [样本数, 2, 128]（2通道：幅度+相位）
    """
    # 确定原始CSI路径
    if scenario == "LOS01_20250817":
        csi_path = os.path.join(LOS_FP_PATH, f"{device_name}.npz")
    elif scenario == "NLOS01_20250821":
        csi_path = os.path.join(NLOS_FP_PATH, f"{device_name}.npz")
    elif scenario == "LOS02_20250817":
        csi_path = os.path.join(LOS2_FP_PATH, f"{device_name}.npz")
    elif scenario == "NLOS02_20250821":
        csi_path = os.path.join(NLOS2_FP_PATH, f"{device_name}.npz")
    elif scenario == "DIRECT01_20250821":
        csi_path = os.path.join(DL_FP_PATH, f"{device_name}.npz")
    elif scenario == "DIRECT02_20250821":
        csi_path = os.path.join(DL2_FP_PATH, f"{device_name}.npz")
    elif scenario == "MOVE01_20250822":
        csi_path = os.path.join(MOV_FP_PATH, f"{device_name}.npz")
    elif scenario == "MOVE02_20250822":
        csi_path = os.path.join(MOV2_FP_PATH, f"{device_name}.npz")
    else:
        raise ValueError(f"场景{scenario}不支持原始CSI加载")

    if not os.path.exists(csi_path):
        raise FileNotFoundError(f"原始CSI文件不存在：{csi_path}")

        # 2. 读取npz文件并解析复数CSI
    try:
        npz_data = np.load(csi_path)
        # 支持两种格式：直接复数数组 或 实部+虚部
        if 'csi_complex' in npz_data:
            csi_complex = npz_data['csi_complex']
            csi_complex = np.squeeze(csi_complex)  # 去除冗余维度
        elif 'real' in npz_data and 'imag' in npz_data:
            real_array = npz_data['real']
            imag_array = npz_data['imag']
            if real_array.shape != imag_array.shape:
                raise ValueError(f"实部与虚部维度不一致：{real_array.shape} vs {imag_array.shape}")
            # 处理维度（确保2D：[样本数, 子载波数]）
            if real_array.ndim == 3:
                real_array = real_array.reshape(real_array.shape[0], -1)
                imag_array = imag_array.reshape(imag_array.shape[0], -1)
            csi_complex = real_array + 1j * imag_array
        else:
            raise KeyError("npz文件必须包含'csi_complex'或'real'+'imag'字段")
    except Exception as e:
        raise RuntimeError(f"读取{device_name}.npz失败：{str(e)}")

        # 3. 标准化子载波数（与配置一致）
    current_samples, current_subcarrier = csi_complex.shape
    target_subcarrier = RAW_CSI_FEAT_DIM  # 100
    if current_subcarrier != target_subcarrier:
        x_old = np.arange(current_subcarrier)
        x_new = np.linspace(0, current_subcarrier - 1, target_subcarrier)
        # 对实部和虚部分别插值
        interp_real = interp1d(x_old, np.real(csi_complex).T, kind="linear", axis=0)
        interp_imag = interp1d(x_old, np.imag(csi_complex).T, kind="linear", axis=0)
        csi_complex = interp_real(x_new).T + 1j * interp_imag(x_new).T

    # 4. 转换为幅度+相位双通道（[样本数, 2, 子载波数]）
    amplitude = np.abs(csi_complex)  # [样本数, 100]
    phase = np.angle(csi_complex)  # [样本数, 100]
    csi_feat = np.stack([amplitude, phase], axis=1)  # [样本数, 2, 100]
    return csi_feat


def load_raw_csi_complex(device_name, scenario="LOS"):
    """
    适配npz格式文件：读取存储CSI数据的npz文件，返回标准化复数CSI
    支持两种存储格式：
    1. 包含'real'和'imag'数组（实部和虚部）
    2. 直接包含'csi_complex'复数数组
    """

    import numpy as np
    from scipy.interpolate import interp1d

    # 确定原始CSI路径
    if scenario == "LOS01_20250817":
        csi_path = os.path.join(LOS_FP_PATH, f"{device_name}.npz")
    elif scenario == "NLOS01_20250821":
        csi_path = os.path.join(NLOS_FP_PATH, f"{device_name}.npz")
    elif scenario == "LOS02_20250817":
        csi_path = os.path.join(LOS2_FP_PATH, f"{device_name}.npz")
    elif scenario == "NLOS02_20250821":
        csi_path = os.path.join(NLOS2_FP_PATH, f"{device_name}.npz")
    elif scenario == "DIRECT01_20250821":
        csi_path = os.path.join(DL_FP_PATH, f"{device_name}.npz")
    elif scenario == "DIRECT02_20250821":
        csi_path = os.path.join(DL2_FP_PATH, f"{device_name}.npz")
    elif scenario == "MOVE01_20250822":
        csi_path = os.path.join(MOV_FP_PATH, f"{device_name}.npz")
    elif scenario == "MOVE02_20250822":
        csi_path = os.path.join(MOV2_FP_PATH, f"{device_name}.npz")
    else:
        raise ValueError(f"场景{scenario}不支持CSI加载")

    # 2. 读取npz文件
    try:
        npz_data = np.load(csi_path)

        # 3. 处理两种数据格式
        if 'csi_complex' in npz_data:
            # 格式1：直接存储复数数组
            csi_complex = npz_data['csi_complex']
            csi_complex = np.squeeze(csi_complex)  # 去除冗余维度
            print(f"[调试] npz复数格式：{device_name} 数据维度{csi_complex.shape}")

        elif 'real' in npz_data and 'imag' in npz_data:
            # 格式2：实部+虚部分开存储
            real_array = npz_data['real']
            imag_array = npz_data['imag']

            # 验证维度一致性
            if real_array.shape != imag_array.shape:
                raise ValueError(f"实部维度{real_array.shape}与虚部维度{imag_array.shape}不一致")

            # 处理维度（确保2D）
            if real_array.ndim == 3:
                real_array = real_array.reshape(real_array.shape[0], -1)
                imag_array = imag_array.reshape(imag_array.shape[0], -1)

            # 重构复数
            csi_complex = real_array + 1j * imag_array
            print(f"[调试] npz实虚部格式：{device_name} 数据维度{csi_complex.shape}")

        else:
            raise KeyError("npz文件必须包含'csi_complex'或'real'+'imag'字段")

    except Exception as e:
        raise RuntimeError(f"读取{device_name}.npz失败：{str(e)}")

    # 4. 标准化子载波数（保持原有逻辑）
    target_subcarrier = RAW_CSI_FEAT_DIM
    current_samples, current_subcarrier = csi_complex.shape

    if current_subcarrier != target_subcarrier:
        print(f"[调试] 子载波数标准化：{current_subcarrier}→{target_subcarrier}")
        x_old = np.arange(current_subcarrier)
        x_new = np.linspace(0, current_subcarrier - 1, target_subcarrier)

        interp_real = interp1d(x_old, np.real(csi_complex).T, kind="linear", axis=0)
        interp_imag = interp1d(x_old, np.imag(csi_complex).T, kind="linear", axis=0)

        csi_complex = interp_real(x_new).T + 1j * interp_imag(x_new).T

    return csi_complex



# tools_1.py 新增代码    实验新增信道集合
def load_dut11_data(scenario="DL"):
    """
    加载DUT11的信道数据（用于剥离环境干扰）
    :param scenario: 场景（DL=理想信道，LOS/NLOS/MOV=带干扰信道）
    :return: DUT11的复数CSI，形状[1000, 128]
    """
    from auth_config import DUT11_FP_PATH  # 需在auth_config.py新增DUT11路径
    # DUT11数据路径：按场景分类，如DUT11_DL.mat、DUT11_LOS.mat
    dut11_path = os.path.join(DUT11_FP_PATH, f"DUT11_{scenario}.mat")
    if not os.path.exists(dut11_path):
        raise FileNotFoundError(f"DUT11 {scenario}场景文件不存在：{dut11_path}")

    from scipy.io import loadmat
    dut11_data = loadmat(dut11_path)['Combined_CSI']  # 假设.mat中变量名为'csi_dut11'
    # 验证维度（1000样本×128子载波，复数）
    if dut11_data.shape != (1000, 128) or not np.iscomplexobj(dut11_data):
        raise ValueError(f"DUT11 {scenario}场景数据错误：需为(1000,128)复数")
    return dut11_data


def generate_universal_channel_noise(epsilon=1e-6):
    """
    用DUT11生成通用信道干扰集合（LOS/NLOS/MOV三种干扰）
    :return: 通用干扰字典，key=场景，value=干扰特征[1000, 128]（复数）
    """
    # 1. 加载DUT11的理想信道（DL）作为基准
    csi_dut11_dl = load_dut11_data(scenario="DL")
    # 2. 加载DUT11的多环境信道（LOS/NLOS/MOV）
    env_scenarios = ["LOS", "NLOS", "MOV"]
    universal_noise = {}

    for env in env_scenarios:
        csi_dut11_env = load_dut11_data(scenario=env)
        # 计算通用干扰：环境信道 ÷ 理想信道（剥离DUT11自身特征，仅剩环境干扰）
        denominator = np.abs(csi_dut11_dl) ** 2 + epsilon
        real_part = (np.real(csi_dut11_env) * np.real(csi_dut11_dl) + np.imag(csi_dut11_env) * np.imag(
            csi_dut11_dl)) / denominator
        imag_part = (np.imag(csi_dut11_env) * np.real(csi_dut11_dl) - np.real(csi_dut11_env) * np.imag(
            csi_dut11_dl)) / denominator
        noise = real_part + 1j * imag_part
        # 归一化干扰特征（避免数值波动）
        noise = noise / (np.abs(noise).max() + epsilon)
        universal_noise[env] = noise
        print(f"生成DUT11 {env}场景通用干扰：{noise.shape}（复数）")
    return universal_noise


import random  # 需导入随机模块

import random
import numpy as np
import torch  # 可学习权重需依赖PyTorch（若用numpy实现可去掉）


def fuse_noise_with_attention(
        target_csi,
        universal_noise,
        attention_mode="fixed",  # "fixed"=固定权重，"learnable"=可学习权重
        w_original=0.8,  # 原始特征固定权重（仅attention_mode="fixed"生效）
        w_noise=0.2,  # 干扰特征固定权重（仅attention_mode="fixed"生效）
        keep_original=True
):
    """
    带注意力机制的特征融合：强化原始特征，弱化信道干扰
    :param target_csi: 目标设备原始CSI [1000, 128]（复数）
    :param universal_noise: 通用干扰字典（LOS/NLOS/MOV）
    :param attention_mode: 注意力模式
    :param w_original: 原始特征固定权重（0~1）
    :param w_noise: 干扰特征固定权重（0~1，需满足w_original + w_noise = 1）
    :return: 注意力加权后的融合数据 [1000, 128]（复数）
    """
    # 1. 随机选择一种干扰（复用之前的随机逻辑）
    if not universal_noise:
        raise ValueError("通用干扰字典为空！需先生成干扰")
    random_scenario = random.choice(list(universal_noise.keys()))
    random_noise = universal_noise[random_scenario]
    # 确保干扰与原始数据样本数一致
    if random_noise.shape[0] != target_csi.shape[0]:
        random_noise = random_noise[:target_csi.shape[0]]

    # 2. 生成干扰样本（目标CSI × 随机干扰）
    noise_sample = target_csi * random_noise  # [1000, 128]（复数）

    # 3. 注意力权重计算（核心步骤）
    if attention_mode == "fixed":
        # 模式1：规则化固定权重（快速验证）
        if not np.isclose(w_original + w_noise, 1.0):
            raise ValueError(f"固定权重之和需为1！当前w_original={w_original}, w_noise={w_noise}")
        # 对原始特征和干扰特征分别加权（复数乘法：权重×复数）
        weighted_original = target_csi * w_original
        weighted_noise = noise_sample * w_noise
        # 融合：加权求和（保留复数特性）
        fused_data = weighted_original + weighted_noise
        print(f"✅ 固定注意力权重：原始特征={w_original}, 干扰特征={w_noise}，干扰场景={random_scenario}")

    elif attention_mode == "learnable":
        # 模式2：可学习注意力权重（让模型自主学习权重分配）
        # 步骤1：将复数特征转换为实部+虚部的双通道特征（适配注意力网络）
        def complex_to_real(complex_data):
            """复数→实部+虚部双通道：[1000, 128] → [1000, 2, 128]"""
            real_part = np.real(complex_data)
            imag_part = np.imag(complex_data)
            return np.stack([real_part, imag_part], axis=1)

        # 转换原始特征和干扰特征为实数值（注意力网络无法直接处理复数）
        original_real = complex_to_real(target_csi)  # [1000, 2, 128]
        noise_real = complex_to_real(noise_sample)  # [1000, 2, 128]

        # 步骤2：定义轻量级注意力网络（输入：特征统计量，输出：2个权重）
        class FeatureAttention(nn.Module):
            def __init__(self, feat_dim=128):
                super().__init__()
                # 输入：原始特征和干扰特征的全局统计量（各2个值：实部均值+虚部均值）
                self.fc = nn.Sequential(
                    nn.Linear(4, 8),  # 4个输入（原始实均、原始虚均、干扰实均、干扰虚均）
                    nn.ReLU(),
                    nn.Linear(8, 2),  # 输出2个权重（原始特征权重、干扰特征权重）
                    nn.Softmax(dim=1)  # 归一化权重（确保和为1）
                )

            def forward(self, original, noise):
                # 计算全局统计量（每个样本的实部/虚部均值，反映特征强度）
                original_real_mean = original.mean(dim=(1, 2)).unsqueeze(1)  # [1000, 1]
                original_imag_mean = original[:, 1, :].mean(dim=1).unsqueeze(1)  # [1000, 1]
                noise_real_mean = noise.mean(dim=(1, 2)).unsqueeze(1)  # [1000, 1]
                noise_imag_mean = noise[:, 1, :].mean(dim=1).unsqueeze(1)  # [1000, 1]

                # 拼接统计量作为注意力输入
                attn_input = torch.cat([original_real_mean, original_imag_mean,
                                        noise_real_mean, noise_imag_mean], dim=1)  # [1000, 4]
                # 输出归一化权重（[1000, 2]：第0列=原始权重，第1列=干扰权重）
                attn_weights = self.fc(attn_input)
                return attn_weights

        # 步骤3：初始化注意力网络并计算权重（可嵌入训练流程，此处为简化示例）
        attn_net = FeatureAttention(feat_dim=128)
        # 转换为张量（若在训练中，需与模型放在同一设备）
        original_tensor = torch.tensor(original_real, dtype=torch.float32)
        noise_tensor = torch.tensor(noise_real, dtype=torch.float32)
        # 计算注意力权重（[1000, 2]）
        attn_weights = attn_net(original_tensor, noise_tensor)  # 训练中需反向传播更新参数
        w_original_tensor = attn_weights[:, 0].unsqueeze(1).unsqueeze(1)  # [1000, 1, 1]
        w_noise_tensor = attn_weights[:, 1].unsqueeze(1).unsqueeze(1)  # [1000, 1, 1]

        # 步骤4：加权融合（实部+虚部分别加权，再重构为复数）
        weighted_original_real = original_tensor * w_original_tensor
        weighted_noise_real = noise_tensor * w_noise_tensor
        fused_real = weighted_original_real + weighted_noise_real  # [1000, 2, 128]

        # 实部+虚部→复数
        fused_real_part = fused_real[:, 0, :].detach().numpy()  # 训练中无需detach
        fused_imag_part = fused_real[:, 1, :].detach().numpy()
        fused_data = fused_real_part + 1j * fused_imag_part  # [1000, 128]（复数）

        # 打印权重统计信息（验证是否强化原始特征）
        avg_w_original = w_original_tensor.mean().item()
        avg_w_noise = w_noise_tensor.mean().item()
        print(
            f"✅ 可学习注意力权重（平均）：原始特征={avg_w_original:.3f}, 干扰特征={avg_w_noise:.3f}，干扰场景={random_scenario}")

    else:
        raise ValueError(f"不支持的注意力模式：{attention_mode}，仅支持'fixed'或'learnable'")

    # 4. 可选：保留原始数据（若需要同时包含“原始+注意力融合”两类样本）
    if keep_original:
        fused_data = np.concatenate([target_csi, fused_data], axis=0)  # [2000, 128]
        print(f"✅ 最终融合数据形状：{fused_data.shape}（原始1000 + 注意力融合1000）")
    else:
        print(f"✅ 最终融合数据形状：{fused_data.shape}（仅注意力融合1000）")

    return fused_data
